#include "utils.h"
#include <cblas64.h>
#include <stdio.h>
// #include <time.h>
#include <sys/time.h>
#include <x86intrin.h>
#include <Eigen/Dense>

#define NN 256 * 20
#define UNROLL 4
#define BLOCKSIZE 32
#define ID(x, y, s) ((x) * (s) + (y))

double *A, *B, *C;

void dgemm(size_t n, double *A, double *B, double *C) {
  for (size_t i = 0; i < n; ++i)
    for (size_t j = 0; j < n; ++j) {
      double cij = C[ID(i, j, n)];
      for (size_t k = 0; k < n; k++)
        cij += A[ID(i, k, n)] * B[ID(k, j, n)];
      C[ID(i, j, n)] = cij;
    }
}

void dgemm_avx(size_t n, double *B, double *A, double *C) {
  for (size_t i = 0; i < n; i += 4) {
    for (size_t j = 0; j < n; j++) {
      __m256d c0 = _mm256_load_pd(C + i + j * n);
      for (size_t k = 0; k < n; k++) {
        c0 = _mm256_add_pd(c0,
                           _mm256_mul_pd(_mm256_load_pd(A + i + k * n),
                                         _mm256_broadcast_sd(B + k + j * n)));
      }
      _mm256_store_pd(C + i + j * n, c0);
      ;
    }
  }
}

void dgemm_avx_unroll(size_t n, double *B, double *A, double *C) {
  for (size_t i = 0; i < n; i += 4 * UNROLL) {
    for (size_t j = 0; j < n; j++) {
      __m256d c[UNROLL];
      for (int x = 0; x < UNROLL; x++) {
        c[x] = _mm256_load_pd(C + i + x * 4 + j * n);
      }
      for (size_t k = 0; k < n; k++) {
        __m256d b = _mm256_broadcast_sd(B + k + j * n);
        for (int x = 0; x < UNROLL; x++) {
          c[x] = _mm256_add_pd(
              c[x], _mm256_mul_pd(_mm256_load_pd(A + i + k * n + x * 4), b));
        }
      }
      for (int x = 0; x < UNROLL; x++) {
        _mm256_store_pd(C + i + x * 4 + j * n, c[x]);
      }
    }
  }
}

static inline void do_block(int n, int si, int sj, int sk, double *A, double *B,
                            double *C) {
  for (int i = si; i < si + BLOCKSIZE; i += UNROLL * 4) {
    for (int j = sj; j < sj + BLOCKSIZE; j++) {
      __m256d c[UNROLL];
      for (int x = 0; x < UNROLL; x++)
        c[x] = _mm256_load_pd(C + i + x * 4 + j * n);

      for (int k = sk; k < sk + BLOCKSIZE; k++) {
        __m256d b = _mm256_broadcast_sd(B + k + j * n);
        for (int x = 0; x < UNROLL; x++)
          c[x] = _mm256_add_pd(
              c[x], _mm256_mul_pd(_mm256_load_pd(A + n * k + x * 4 + i), b));
      }

      for (int x = 0; x < UNROLL; x++)
        _mm256_store_pd(C + i + x * 4 + j * n, c[x]);
    }
  }
}

void dgemm_avx_unroll_block(size_t n, double *B, double *A, double *C) {
  for (size_t i = 0; i < n; i += BLOCKSIZE)
    for (size_t j = 0; j < n; j += BLOCKSIZE)
      for (size_t k = 0; k < n; k += BLOCKSIZE)
        do_block(n, i, j, k, A, B, C);
}

void dgemm_avx_unroll_block_mp(size_t n, double *B, double *A, double *C) {
#pragma omp parallel for
  for (size_t i = 0; i < n; i += BLOCKSIZE)
    for (size_t j = 0; j < n; j += BLOCKSIZE)
      for (size_t k = 0; k < n; k += BLOCKSIZE)
        do_block(n, i, j, k, A, B, C);
}

void dgemm_blas(size_t n, double *A, double *B, double *C) {
  cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, n, n, n, 1, A, n, B, n,
              0, C, n);
}

void dgemm_eigen(size_t n, double *A, double *B, double *C) {
  using namespace Eigen;
  typedef Matrix<double, Dynamic, Dynamic,RowMajor>rMatrixXd;
  rMatrixXd tmp = (Map<rMatrixXd>(A, n, n) * Map<rMatrixXd>(B, n, n));
  memcpy(C, tmp.data(), n * n * sizeof(double));
}

void test(size_t n, double *A, double *B, double *C,
          void (*dgemm)(size_t, double *, double *, double *)) {
  struct timeval t_start, t_end;
  long timeuse = 0;

  init(&A, &B, &C, NN);
  gettimeofday(&t_start, NULL);
  dgemm(NN, A, B, C);
  gettimeofday(&t_end, NULL);
  timeuse = 1000000 * (t_end.tv_sec - t_start.tv_sec) + t_end.tv_usec -
            t_start.tv_usec;
  printf("time=%f\n", timeuse / 1000000.0);

  show(C, NN);
  mfree(A, B, C);
}

int main() {
  // test(NN, A, B, C, dgemm);
  // test(NN, A, B, C, dgemm_avx);
  // test(NN, A, B, C, dgemm_avx_unroll);
  // test(NN, A, B, C, dgemm_avx_unroll_block);
  test(NN, A, B, C, dgemm_avx_unroll_block_mp);
  test(NN, A, B, C, dgemm_blas);
  test(NN, A, B, C, dgemm_eigen);
}